module dSFMT_interface

! copyright James Spencer 2012.
! New BSD License, see License.txt for details.

! This contains a wrapper around a subset of of the functionality in dSFMT RNG
! (double precision SIMD-oriented Fast Mersenne Twister random number
! generator).

! Functions which return a single (pseudo-)random number are defined as inline
! in the original dSFMT source.  Accessing these functions is quite slow
! (compared to C) in Fortran as we can't inline them (or indeed call them
! without modifying the C source).  A more efficient way is to fill up an array
! of random numbers and then access the array as a stream, refilling the array
! as needed.  By thus amortizing the cost of calling non-inlined functions
! across many random numbers, the cost becomes very competitive with
! a C implementation.  Functions are provided which return a single random
! number and refill the storage array when required, thus giving a fast
! interface which closely mirrors that of the C implementation.

! Usage (for each psuedo-random stream):
!   i)   call dSFMT_init.
!   ii)  access desired number of random numbers either by calling
!        a get_rand_xxx function the required number of times or by filling up
!        user-defined arrays with a fill_array_xxx subroutine.
!   iii) call dSFMT_end to free all memory.

! Note that an array must have at least a minimum number of elements (as given
! in dsfmt_get_min_array_size).

! IMPORTANT: see the warning below about calling multiple get_rand_xxx
! functions.

! Match definition of DSFMT_MEXP in dSFMT.h.
#if !defined(DSFMT_MEXP)
#ifdef __GNUC__
#warning "DSFMT_MEXP is not defined. I assume DSFMT_MEXP is 19937."
#endif
#define DSFMT_MEXP 19937
#endif

use, intrinsic :: iso_c_binding
implicit none

private
public :: dSFMT_t, dSFMT_init, dSFMT_end, dSFMT_reset,  &
          fill_array_close_open, &
          fill_array_open_close, &
          fill_array_open_open,  &
          fill_array_gaussian, &
          get_rand_close_open, &
          get_rand_open_close, &
          get_rand_open_open,  &
          get_rand_gaussian,   &
          get_rand_arr_close_open,                      &
          get_rand_arr_open_close,                      &
          get_rand_arr_open_open,                       &
          get_rand_arr_gaussian,                        &
          dsfmt_get_min_array_size,                     &
          unset, close_open, open_close, open_open

! Expose functions from C as needed.
! See dSFMT documentation for details.
interface
    pure function dsfmt_get_min_array_size() result(min_size) bind(c, name='dsfmt_get_min_array_size')
        import :: c_int32_t
        integer(c_int32_t) :: min_size
    end function dsfmt_get_min_array_size

    pure function malloc_dsfmt_t() result(dsfmt_state) bind(c, name='malloc_dsfmt_t')
        import :: c_ptr
        type(c_ptr) :: dsfmt_state
    end function malloc_dsfmt_t
    pure subroutine free_dsfmt_t(dsfmt_state) bind(c, name='free_dsfmt_t')
        import :: c_ptr
        type(c_ptr), value, intent(in) :: dsfmt_state
    end subroutine free_dsfmt_t

    pure subroutine dsfmt_chk_init_gen_rand(dsfmt_state, seed, mexp) bind(c, name='dsfmt_chk_init_gen_rand')
        import :: c_ptr, c_int32_t
        type(c_ptr), value, intent(in) :: dsfmt_state
        integer(c_int32_t), value, intent(in) :: seed, mexp
    end subroutine dsfmt_chk_init_gen_rand

    pure subroutine dsfmt_fill_array_close_open(dSFMT_state, array, array_size) bind(c, name='dsfmt_fill_array_close_open')
        import :: c_ptr, c_double, c_int32_t
        type(c_ptr), value, intent(in) :: dSFMT_state
        real(c_double), intent(out) :: array(*)
        integer(c_int32_t), value, intent(in) :: array_size
    end subroutine dsfmt_fill_array_close_open
    pure subroutine dsfmt_fill_array_open_close(dSFMT_state, array, array_size) bind(c, name='dsfmt_fill_array_open_close')
        import :: c_ptr, c_double, c_int32_t
        type(c_ptr), value, intent(in) :: dSFMT_state
        real(c_double), intent(out) :: array(*)
        integer(c_int32_t), value, intent(in) :: array_size
    end subroutine dsfmt_fill_array_open_close
    pure subroutine dsfmt_fill_array_open_open(dSFMT_state, array, array_size) bind(c, name='dsfmt_fill_array_open_open')
        import :: c_ptr, c_double, c_int32_t
        type(c_ptr), value, intent(in) :: dSFMT_state
        real(c_double), intent(out) :: array(*)
        integer(c_int32_t), value, intent(in) :: array_size
    end subroutine dsfmt_fill_array_open_open
end interface

! distribution modes
integer, parameter :: unset = 2**0
integer, parameter :: close_open = 2**1
integer, parameter :: open_close = 2**2
integer, parameter :: open_open = 2**3
integer, parameter :: gaussian = 2**4

type dSFMT_t
    private
    ! Type of distribution in the random store
    integer, public :: distribution
    ! Pointer to the dsfmt_t internal state (as defined in dSFMT.h).
    type(c_ptr) :: dSFMT_state = c_null_ptr
    ! Testing indicates that 50000 is a very good size for the array storing the
    ! random numbers.  Testing was done standalone, so undoubtedly influenced by
    ! cache size and this might be different for real-world applications.
    integer(c_int) :: random_store_size
    ! Seed passed to dSFMT.
    integer :: seed
    ! Store of random numbers.
    real(c_double), allocatable :: random_store(:)
    ! Index of the next random number to be returned from random_store.
    integer :: next_element
end type dSFMT_t

integer, parameter, private :: dp = selected_real_kind(15,307)

contains

!--- Initialisation, termination and state interaction ---

    subroutine dSFMT_init(seed, rng_store_size, rng)

        ! Initialise the dSFMT RNG and fill rng%random_store with
        ! a block of random numbers in interval [0,1).
        !
        ! In:
        !    seed: seed for the RNG.
        !    rng_store_size: number of random numbers to store at once.
        !       dsfmt_get_min_array_size() is used if
        !       rng_store_size < dsfmt_get_min_array_size().
        ! Out:
        !    rng: dSFMT_t with internal variables initialised and associated
        !       with an initialised psuedo-random number stream.
        !       If rng has already been initialised, then no additional memory
        !       allocations are performed unless rng_store_size has changed.

        integer, intent(in) :: seed, rng_store_size
        type(dSFMT_t), intent(inout) :: rng

        if (dsfmt_get_min_array_size() > rng_store_size) then
            rng%random_store_size = dsfmt_get_min_array_size()
        else
            rng%random_store_size = rng_store_size
        end if

        if (allocated(rng%random_store)) then
            if (size(rng%random_store) /= rng%random_store_size) then
                deallocate(rng%random_store)
                allocate(rng%random_store(rng%random_store_size))
            end if
        else
            allocate(rng%random_store(rng%random_store_size))
        end if

        ! Set current element to be larger than the store, so it is
        ! filled in the first call to the get_* functions.
        rng%next_element = rng%random_store_size + 1
        ! But the store has not yet been filled.
        rng%distribution = unset

        rng%seed = int(seed, c_int32_t)

        ! Create dSFMT state
        ! Why is c_associated not pure?  I have not been able to find any good
        ! reason for this to be the case...
        if (.not.c_associated(rng%dSFMT_state)) rng%dSFMT_state = malloc_dsfmt_t()

        ! Initialise dSFMT PRNG.
        call dsfmt_chk_init_gen_rand(rng%dSFMT_state, rng%seed, DSFMT_MEXP)

    end subroutine dSFMT_init

    pure subroutine dSFMT_end(rng)

        ! Deallocate and reset a dSFMT_t variable safely.

        ! In/Out:
        !    rng: dSFMT_t.  The dSFMT state is closed on output and memory
        !       associated with it is deallocated.

        type(dSFMT_t), intent(inout) :: rng

        ! Destroy dSFMT state
        call free_dsfmt_t(rng%dSFMT_state)

        ! Reset store.
        deallocate(rng%random_store)
        rng%random_store_size = -1
        rng%next_element = -1

    end subroutine dSFMT_end

    pure subroutine dSFMT_reset(rng)

        ! Reset the dSFMT state such that the next get_rand_xxx call requires
        ! the store of random numbers to be refilled.

        ! In/Out:
        !    rng: dSFMT_t.  The internal store of random numbers is cleared.

        type(dSFMT_t), intent(inout) :: rng

        rng%next_element = rng%random_store_size + 1
        rng%random_store = huge(1.0_dp)
        rng%distribution = unset

    end subroutine dSFMT_reset

!--- Get arrays of random numbers ---

    pure subroutine fill_array_close_open(rng, array)

        ! In/Out:
        !    rng: dSFMT_t.  The dSFMT state is updated by the request for random
        !       numbers.
        ! Out:
        !    array: array filled with random numbers in the interval [0,1).

        ! NOTE: array must be at least as large as the size returned by dsfmt_get_min_array_size.

        type(dSFMT_t), intent(inout) :: rng
        real(c_double), intent(out) :: array(:)

        call dsfmt_fill_array_close_open(rng%dSFMT_state, array, size(array))

    end subroutine fill_array_close_open

    pure subroutine fill_array_open_close(rng, array)

        ! In/Out:
        !    rng: dSFMT_t.  The dSFMT state is updated by the request for random
        !       numbers.
        ! Out:
        !    array: array filled with random numbers in the interval (0,1].

        ! NOTE: array must be at least as large as the size returned by dsfmt_get_min_array_size.

        type(dSFMT_t), intent(inout) :: rng
        real(c_double), intent(out) :: array(:)

        call dsfmt_fill_array_open_close(rng%dSFMT_state, array, size(array))

    end subroutine fill_array_open_close

    pure subroutine fill_array_open_open(rng, array)

        ! In/Out:
        !    rng: dSFMT_t.  The dSFMT state is updated by the request for random
        !       numbers.
        ! Out:
        !    array: array filled with random numbers in the interval (0,1).

        ! NOTE: array must be at least as large as the size returned by dsfmt_get_min_array_size.

        type(dSFMT_t), intent(inout) :: rng
        real(c_double), intent(out) :: array(:)

        call dsfmt_fill_array_open_open(rng%dSFMT_state, array, size(array))

    end subroutine fill_array_open_open

    pure subroutine fill_array_gaussian(rng, array)

        ! In/Out:
        !    rng: dSFMT_t.  The dSFMT state is updated by the request for random
        !       numbers.
        ! Out:
        !    array: array filled with random numbers with a standard normal
        !    (Gaussian) distribution (i.e. mean=0, variance=1).

        ! NOTE: array must be at least as large as the size returned by dsfmt_get_min_array_size.

        type(dSFMT_t), intent(inout) :: rng
        real(c_double), intent(out) :: array(:)

        real(dp), parameter :: pi = acos(-1.0_dp)
        integer :: i
        real(c_double) :: s, t

        call dsfmt_fill_array_open_close(rng%dSFMT_state, array, size(array))

        ! The Box-Muller transform converts a pair of random numbers, u, v,
        ! which are uniformly distributed in (0,1] into a pair of random
        ! numbers, y, z, which have a standard normal distribution using
        !    y = R cos(\phi) = \sqrt(-2*ln(u)) cos(2\pi v)
        !    z = R sin(\phi) = \sqrt(-2*ln(u)) sin(2\pi v)
        ! See http://en.wikipedia.org/wiki/Box_Muller_transform for more
        ! details.
        ! This might be slower than the Marsaglia polar method but it doesn't
        ! involve any rejections, which is convenient for us.
        do i = 1, ubound(array, dim=1), 2
            s = sqrt(-2*log(array(i)))
            t = 2*pi*array(i+1)
            array(i) = s*cos(t)
            array(i+1) = s*sin(t)
        end do

    end subroutine fill_array_gaussian

!--- Get a random number ---

! WARNING: these calls use a store internal to the dSFMT_t variable.  Calls to
! different functions with the same state must be separated either by a call to
! dSFMT_reset or a call to dSFMT_end followed by reinitialisation of the dSFMT_t
! state using dSFMT_init.

    function get_rand_close_open(rng) result(r)

        ! In/Out:
        !    rng: dSFMT_t.  The dSFMT state is updated by the request for random
        !       numbers.  The store of random numbers is refilled if necessary.

        ! Returns:
        !    random number in interval [0,1).

        real(dp) :: r
        type(dSFMT_t), intent(inout) :: rng

        if (rng%next_element == rng%random_store_size+1) then
            ! Run out of random numbers: get more.
            call dsfmt_fill_array_close_open(rng%dSFMT_state, rng%random_store, rng%random_store_size)
            rng%next_element = 1
            rng%distribution = close_open
        end if

        r = rng%random_store(rng%next_element)
        rng%next_element = rng%next_element + 1 

    end function get_rand_close_open

    function get_rand_open_close(rng) result(r)

        ! In/Out:
        !    rng: dSFMT_t.  The dSFMT state is updated by the request for random
        !       numbers.  The store of random numbers is refilled if necessary.

        ! Returns:
        !    random number in interval (0,1].

        real(dp) :: r
        type(dSFMT_t), intent(inout) :: rng

        if (rng%next_element == rng%random_store_size+1) then
            ! Run out of random numbers: get more.
            call dsfmt_fill_array_open_close(rng%dSFMT_state, rng%random_store, rng%random_store_size)
            rng%next_element = 1
            rng%distribution = open_close
        end if

        r = rng%random_store(rng%next_element)
        rng%next_element = rng%next_element + 1 

    end function get_rand_open_close

    function get_rand_open_open(rng) result(r)

        ! In/Out:
        !    rng: dSFMT_t.  The dSFMT state is updated by the request for random
        !       numbers.  The store of random numbers is refilled if necessary.

        ! Returns:
        !    random number in interval (0,1).

        real(dp) :: r
        type(dSFMT_t), intent(inout) :: rng

        if (rng%next_element == rng%random_store_size+1) then
            ! Run out of random numbers: get more.
            call dsfmt_fill_array_open_open(rng%dSFMT_state, rng%random_store, rng%random_store_size)
            rng%next_element = 1
            rng%distribution = open_open
        end if

        r = rng%random_store(rng%next_element)
        rng%next_element = rng%next_element + 1 

    end function get_rand_open_open

    function get_rand_gaussian(rng) result(r)

        ! In/Out:
        !    rng: dSFMT_t.  The dSFMT state is updated by the request for random
        !       numbers.  The store of random numbers is refilled if necessary.

        ! Returns:
        !    random number with a standard normal (Gaussian) distribution (ie
        !    mean=0 and variance=1).

        real(dp) :: r
        type(dSFMT_t), intent(inout) :: rng

        if (rng%next_element == rng%random_store_size+1) then
            ! Run out of random numbers: get more.
            call fill_array_gaussian(rng, rng%random_store)
            rng%next_element = 1
            rng%distribution = gaussian
        end if

        r = rng%random_store(rng%next_element)
        rng%next_element = rng%next_element + 1

    end function get_rand_gaussian

!--- Get array of random numbers from a store ---

! If a small number of random numbers is repeatedly required, then filling up
! an array at a time comes with a large function call overhead.  The following
! routines use an internal store, similar to the get_rand_xxx functions, but
! fill a small array in a single call.  This can be faster than repeated
! get_rand_xxx calls as we need to only test how many elements are left in the
! store once rather than N times.

! WARNING: these calls use a store internal to the dSFMT_t variable.  Calls to
! different functions with the same state must be separated either by a call to
! dSFMT_reset or a call to dSFMT_end followed by reinitialisation of the dSFMT_t
! state using dSFMT_init.

    pure subroutine get_rand_arr_close_open(rng, arr, n)

        ! Fill an array with random numbers in interval [0,1).

        ! In/Out:
        !    rng: dSFMT_t.  The dSFMT state is updated by the request for random
        !       numbers.  The store of random numbers is refilled if necessary.
        ! Out:
        !    arr: array of random numbers.  Must contain at least n elements.
        ! In:
        !    n: number of random numbers to return in arr.

        type(dSFMT_t), intent(inout) :: rng
        real(dp), intent(out) :: arr(:)
        integer, intent(in) :: n

        integer :: navail, nleft

        if (rng%next_element + n <= rng%random_store_size) then
            arr(1:n) = rng%random_store(rng%next_element:rng%next_element+n-1)
            rng%next_element = rng%next_element + n
        else
            navail = rng%random_store_size - rng%next_element + 1
            arr(1:navail) = rng%random_store(rng%next_element:rng%random_store_size)
            call dsfmt_fill_array_close_open(rng%dSFMT_state, rng%random_store, rng%random_store_size)
            rng%distribution = close_open
            nleft = n - navail
            arr(navail+1:n) = rng%random_store(1:nleft)
            rng%next_element = nleft + 1
        end if

    end subroutine get_rand_arr_close_open

    pure subroutine get_rand_arr_open_close(rng, arr, n)

        ! Fill an array with random numbers in interval (0,1].

        ! In/Out:
        !    rng: dSFMT_t.  The dSFMT state is updated by the request for random
        !       numbers.  The store of random numbers is refilled if necessary.
        ! Out:
        !    arr: array of random numbers.  Must contain at least n elements.
        ! In:
        !    n: number of random numbers to return in arr.

        type(dSFMT_t), intent(inout) :: rng
        real(dp), intent(out) :: arr(:)
        integer, intent(in) :: n

        integer :: navail, nleft

        if (rng%next_element + n <= rng%random_store_size) then
            arr(1:n) = rng%random_store(rng%next_element:rng%next_element+n-1)
            rng%next_element = rng%next_element + n
        else
            navail = rng%random_store_size - rng%next_element + 1
            arr(1:navail) = rng%random_store(rng%next_element:rng%random_store_size)
            call dsfmt_fill_array_open_close(rng%dSFMT_state, rng%random_store, rng%random_store_size)
            rng%distribution = open_close
            nleft = n - navail
            arr(navail+1:n) = rng%random_store(1:nleft)
            rng%next_element = nleft + 1
        end if

    end subroutine get_rand_arr_open_close

    pure subroutine get_rand_arr_open_open(rng, arr, n)

        ! Fill an array with random numbers in interval (0,1).

        ! In/Out:
        !    rng: dSFMT_t.  The dSFMT state is updated by the request for random
        !       numbers.  The store of random numbers is refilled if necessary.
        ! Out:
        !    arr: array of random numbers.  Must contain at least n elements.
        ! In:
        !    n: number of random numbers to return in arr.

        type(dSFMT_t), intent(inout) :: rng
        real(dp), intent(out) :: arr(:)
        integer, intent(in) :: n

        integer :: navail, nleft

        if (rng%next_element + n <= rng%random_store_size) then
            arr(1:n) = rng%random_store(rng%next_element:rng%next_element+n-1)
            rng%next_element = rng%next_element + n
        else
            navail = rng%random_store_size - rng%next_element + 1
            arr(1:navail) = rng%random_store(rng%next_element:rng%random_store_size)
            call dsfmt_fill_array_open_open(rng%dSFMT_state, rng%random_store, rng%random_store_size)
            rng%distribution = open_open
            nleft = n - navail
            arr(navail+1:n) = rng%random_store(1:nleft)
            rng%next_element = nleft + 1
        end if

    end subroutine get_rand_arr_open_open

    pure subroutine get_rand_arr_gaussian(rng, arr, n)

        ! Fill an array with random numbers in interval (0,1).

        ! In/Out:
        !    rng: dSFMT_t.  The dSFMT state is updated by the request for random
        !       numbers.  The store of random numbers is refilled if necessary.
        ! Out:
        !    arr: array of random numbers.  Must contain at least n elements.
        ! In:
        !    n: number of random numbers to return in arr.

        type(dSFMT_t), intent(inout) :: rng
        real(dp), intent(out) :: arr(:)
        integer, intent(in) :: n

        integer :: navail, nleft

        if (rng%next_element + n <= rng%random_store_size) then
            arr(1:n) = rng%random_store(rng%next_element:rng%next_element+n-1)
            rng%next_element = rng%next_element + n
        else
            navail = rng%random_store_size - rng%next_element + 1
            arr(1:navail) = rng%random_store(rng%next_element:rng%random_store_size)
            call fill_array_gaussian(rng, rng%random_store)
            rng%distribution = gaussian
            nleft = n - navail
            arr(navail+1:n) = rng%random_store(1:nleft)
            rng%next_element = nleft + 1
        end if

    end subroutine get_rand_arr_gaussian

end module dSFMT_interface
